from imagedownloader import yandexapi
from CLIP import clip_score as clip
from ipadapter import IpAdapter as ip
from ip2p import edit_cli as ip2p
from gpt import GPT
import os
import logging
import shutil
import re

if __name__ == '__main__':
    # 控制整体的生成图片顺序
    progress = 1

    # 设置日志级别: INFO
    logging.basicConfig(format='%(asctime)s - %(pathname)s[line:%(lineno)d] - %(levelname)s: %(message)s',
                        level=logging.INFO)

    # 输入文本
    logging.info("===Start===")
    input_text = "Husky wearing a straw hat on the snow"
    logging.info("Input_text: %s", input_text)

    # 保存图片的目录
    img_dir = "./images"
    logging.info("检索图片保存在：%s", img_dir)

    # 原文本检索
    yandexapi.search_image(input_text, img_dir, 30, progress)

    # 调用GPT获得丰富检索词
    query = GPT.get_search_keywords(input_text)

    # 检索策略
    # for i in range(len(query)):
    #     yandexapi.search_image(query[i], img_dir, 5, 2)

    # 筛选器模块
    # CLIP评分模块
    prefix = "a photo of "
    clip_text = prefix + input_text
    pattern = re.compile(r'^image_(\d+)_\d+$', re.IGNORECASE)
    count = 0
    scores = {}
    for image in os.listdir(img_dir):
        basename, ext = os.path.splitext(image)
        match = pattern.match(basename)
        if int(match.group(1)) != progress:
            continue
        image_path = os.path.join(img_dir, image)
        try:
            score = clip.get_clip_score(image_path, clip_text)
            count += 1
        except Exception as e:
            logging.error(f"计算图片{image}的CLIP分数时出错：{e}")
            continue
        scores[image] = score
        logging.info(f"{image} CLIP 分数为：{score}")
    sort_images = sorted(scores, key=scores.get, reverse=True)[:count]
    logging.info("CLIP 分数排序：")
    print(sort_images)
    clip_top_image = sort_images[0]
    logging.info(f"CLIP 评分最高的图片: {clip_top_image}")

    # LLM筛选器
    for image in sort_images:
        image_path = os.path.join(img_dir, image)
        if GPT.filter_image(image_path, input_text):
            clip_top_image = image
            print(f"筛选图片，最终编辑图片: {clip_top_image}")
            break

    # 存储图片
    backups_origin_dir = '/root/data/origins'  # 原始图片
    backups_output_dir = '/root/data/outputs'  # 生成图片
    # 创建文件夹
    if not os.path.exists(backups_origin_dir):
        os.makedirs(backups_origin_dir)
    if not os.path.exists(backups_output_dir):
        os.makedirs(backups_output_dir)
    source_path = os.path.join(img_dir, clip_top_image)
    name, ext = os.path.splitext(clip_top_image)
    backups_origin_dest_path = os.path.join(backups_origin_dir, f'{progress}' + ext)
    logging.info(f"原始图片： {backups_origin_dest_path}")
    backups_output_dest_path = os.path.join(backups_output_dir, f'{progress}' + ext)
    logging.info(f"即将编辑： {backups_output_dest_path}")
    shutil.copy(source_path, backups_origin_dest_path)
    shutil.copy(source_path, backups_output_dest_path)

    # 多阶段编辑与循环反馈模块
    # 编辑指令生成
    instruction_first = GPT.get_instruction_first(backups_output_dest_path, backups_output_dest_path, input_text)

    # 调用编辑模型推理
    if len(instruction_first) != 0 and instruction_first[0] != 'No editing required.' and instruction_first[0] != 'no editing required.':
        for i in range(len(instruction_first)):
            target_image = ip.ip_adapter_sd_inference(backups_output_dest_path, instruction_first[i], 1)
            # sdxl
            # target_image = ip.ip_adapter_sdxl_inference(backups_output_dest_path, instruction_first[i], 1)
            target_image[0].save(backups_output_dest_path)

    # 编辑指令生成
    instruction_second = GPT.get_instruction_second(backups_output_dest_path, backups_output_dest_path, input_text)

    # 调用编辑模型推理
    if len(instruction_second) != 0 and instruction_second[0] != 'No editing required.' and instruction_second[0] != 'no editing required.':
        for i in range(len(instruction_second)):
            target_image = ip2p.run_model(input_path=backups_output_dest_path,
                                          output_path=backups_output_dest_path,
                                          edit_text=instruction_second[i],
                                          ckpt="/root/data/checkpoints/MagicBrush-epoch-52-step-4999.ckpt",
                                          cfg_text=7.5,
                                          cfg_image=1.5,
                                          seed=2000)
            target_image[0].save(backups_output_dest_path)

    # 调用dall-e进行推理
    GPT.generate_image_dall_e3(input_text, f"./dall_e3/{progress}", f"/root/data/dall_e3/{progress}")

